package org.jetlinks.pro.gateway.session;

import lombok.Generated;
import lombok.Getter;
import lombok.Setter;
import org.hswebframework.ezorm.rdb.mapping.ReactiveRepository;
import org.jetlinks.core.cache.Caches;
import org.jetlinks.core.device.DeviceConfigKey;
import org.jetlinks.core.device.DeviceRegistry;
import org.jetlinks.core.device.DeviceState;
import org.jetlinks.core.message.codec.Transport;
import org.jetlinks.core.server.monitor.GatewayServerMonitor;
import org.jetlinks.core.server.session.*;
import org.jetlinks.pro.gateway.session.entity.PersistentSessionEntity;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;
import reactor.core.Disposable;
import reactor.core.Disposables;
import reactor.core.publisher.*;
import reactor.core.scheduler.Schedulers;

import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

/**
 * @author zhouhao
 * @since 1.0.0
 */
public class DefaultDeviceSessionManager implements DeviceSessionManager {

    private final Map<String, DeviceSession> repository = Caches.newCache();

    private final Map<String, Map<String, ChildrenDeviceSession>> children = Caches.newCache();

    @Getter
    @Setter
    private Logger log = LoggerFactory.getLogger(DefaultDeviceSessionManager.class);

    @Getter
    @Setter
    private GatewayServerMonitor gatewayServerMonitor;

    @Getter
    @Setter
    private DeviceRegistry registry;

    @Setter
    private ReactiveRepository<PersistentSessionEntity, String> sessionRepository;

    private final EmitterProcessor<DeviceSession> onDeviceRegister = EmitterProcessor.create(false);
    private final EmitterProcessor<DeviceSession> onDeviceUnRegister = EmitterProcessor.create(false);
    private final EmitterProcessor<DeviceSession> unregisterHandler = EmitterProcessor.create(false);

    private final FluxSink<DeviceSession> unregisterListener = onDeviceUnRegister.sink(FluxSink.OverflowStrategy.BUFFER);
    private final FluxSink<DeviceSession> registerListener = onDeviceRegister.sink(FluxSink.OverflowStrategy.BUFFER);
    private final FluxSink<DeviceSession> unregisterSession = unregisterHandler.sink(FluxSink.OverflowStrategy.BUFFER);

    private final Disposable.Composite disposable = Disposables.composite();

    private String serverId;

    private final Queue<Runnable> scheduleJobQueue = new ArrayDeque<>();

    private final Map<String, LongAdder> transportCounter = new ConcurrentHashMap<>();

    @Getter
    @Setter
    private Map<String, Long> transportLimits = new ConcurrentHashMap<>();

    @Generated
    public void setTransportLimit(Transport transport, long limit) {
        transportLimits.put(transport.getId(), limit);
    }

    public void shutdown() {
        repository.values()
                  .stream()
                  //只离线未持久化的session
                  .filter(session -> !(session instanceof PersistentSession))
                  .map(DeviceSession::getId)
                  .forEach(this::unregister);
        onDeviceRegister.onComplete();
        onDeviceUnRegister.onComplete();
        unregisterHandler.onComplete();
        disposable.dispose();
    }

    @Override
    @Generated
    public boolean isOutOfMaximumSessionLimit(Transport transport) {
        long max = getMaximumSession(transport);
        return max > 0 && getCurrentSession(transport) >= max;
    }

    @Override
    @Generated
    public long getMaximumSession(Transport transport) {
        Long counter = transportLimits.get(transport.getId());
        return counter == null ? -1 : counter;
    }

    @Override
    @Generated
    public long getCurrentSession(Transport transport) {
        LongAdder counter = transportCounter.get(transport.getId());
        return counter == null ? 0 : counter.longValue();
    }

    public Mono<Long> checkSession() {
        AtomicLong startWith = new AtomicLong();
        return Flux
            .fromIterable(repository.values())
            .distinct()
            .publishOn(Schedulers.parallel())
            .filterWhen(session -> {
                if (!session.isAlive() || session.getOperator() == null) {
                    return Mono.just(true);
                }
                return Mono
                    .zip(
                        //状态为unknown可能是已经注销了
                        session.getOperator().getState().filter(state -> state != DeviceState.unknown),
                        session.getOperator().getConnectionServerId().defaultIfEmpty("")
                    )
                    .flatMap((tp2) -> {
                        //serverId不一致
                        if (StringUtils.hasText(serverId) && !serverId.equals(tp2.getT2())) {
                            //设备已经在线,在其他服务器上?
                            if (tp2.getT1().equals(DeviceState.online)) {
                                //检查状态
                                return session
                                    .getOperator()
                                    .checkState()//检查真实状态判断是否真的在线
                                    .flatMap(realState -> {
                                        //确实在其他服务器上
                                        if (realState.equals(DeviceState.online)) {
                                            log.info("device {} connected another server:{},close session:{}",
                                                     session.getDeviceId(),
                                                     tp2.getT2(),
                                                     session);
                                            //注销本地的会话,以后设置的serverId为准
                                            unregister(session.getId(), false);
                                            return Mono.just(false);
                                        }
                                        //不在线?同步缓存信息
                                        return session
                                            .getOperator()
                                            .online(serverId, session.getId(), session
                                                .getClientAddress()
                                                .map(InetSocketAddress::toString)
                                                .orElse(null))
                                            .thenReturn(false);
                                    });
                            } else {
                                //更新设备状态为在线
                                return session
                                    .getOperator()
                                    .online(serverId, session.getId(), session
                                        .getClientAddress()
                                        .map(InetSocketAddress::toString)
                                        .orElse(null))
//                                    .then(Mono.fromRunnable(() -> registerListener.next(session)))
                                    .thenReturn(false);
                            }
                        }
                        return Mono.just(false);
                    })
                    //已经注销的设备
                    .switchIfEmpty(Mono.fromSupplier(() -> {
                        unregister(session.getId(), false);
                        return false;
                    }));
            })
            .map(DeviceSession::getId)
            .doOnNext(this::unregister)
            .count()
            .doOnNext((l) -> {
                if (log.isInfoEnabled() && l > 0) {
                    log.info("expired sessions:{}", l);
                }
            })
            .doOnError(err -> log.error(err.getMessage(), err))
            .doOnSubscribe(subscription -> {
                log.trace("start check session");
                startWith.set(System.currentTimeMillis());
            })
            .doFinally(s -> {
                //上报session数量
                transportCounter.forEach((transport, number) -> gatewayServerMonitor
                    .metrics()
                    .reportSession(transport, number.intValue()));
                //执行任务
                for (Runnable runnable = scheduleJobQueue.poll(); runnable != null; runnable = scheduleJobQueue.poll()) {
                    try {
                        runnable.run();
                    } catch (Exception e) {
                        log.error(e.getMessage(), e);
                    }
                }
                if (log.isTraceEnabled()) {
                    log.trace("check session complete,current server sessions:{}.use time:{}ms.",
                              transportCounter,
                              System.currentTimeMillis() - startWith.get());
                }
            });
    }

    public void init() {
        Objects.requireNonNull(gatewayServerMonitor, "gatewayServerMonitor");
        Objects.requireNonNull(registry, "registry");
        serverId = gatewayServerMonitor.getCurrentServerId();
        disposable
            .add(Flux
                     .interval(Duration.ofSeconds(10), Duration.ofSeconds(30), Schedulers.newSingle("device-session-checker"))
                     .flatMap(i -> this
                         .checkSession()
                         .then(Mono.defer(this::tryPersistent))
                         .onErrorResume((err) -> {
                             log.error(err.getMessage(), err);
                             return Mono.empty();
                         }))
                     .subscribe());


        unregisterHandler
            .publishOn(Schedulers.parallel())
            .flatMap(session -> {
                //注册中心下线
                return session
                    .getOperator()
                    .offline()
                    .doFinally(s -> {
                        if (onDeviceUnRegister.hasDownstreams()) {
                            unregisterListener.next(session);
                        }
                    })
                    .then(this.closeChildSession(session.getDeviceId()))
                    .onErrorResume((err) -> {
                        log.error(err.getMessage(), err);
                        return Mono.empty();
                    });
            })
            .subscribe();

        if (sessionRepository != null) {
            //会话持久化
            onDeviceRegister
                .filter(session -> session instanceof PersistentSession)
                .cast(PersistentSession.class)
                .windowTimeout(200, Duration.ofSeconds(3))
                .flatMap(this::tryPersistent)
                .subscribe();

            //删除会话
            onDeviceUnRegister
                .filter(session -> session instanceof PersistentSession)
                .map(DeviceSession::getId)
                .bufferTimeout(200, Duration.ofSeconds(3))
                .flatMap(this::cleanSession)
                .subscribe();

            //恢复会话
            disposable.add(Mono
                               .delay(Duration.ofSeconds(10))
                               .then(
                                   sessionRepository
                                       .createQuery()
                                       .where(PersistentSessionEntity::getServerId, serverId)
                                       .fetch()
                                       .flatMap(this::resumeSession)
                                       .then()
                               )
                               .subscribe()
            );
        }
    }

    Mono<Void> closeChildSession(String parentDeviceId) {
        return Mono
            .justOrEmpty(children.remove(parentDeviceId))
            .flatMapIterable(Map::values)
            .flatMap(childrenDeviceSession -> childrenDeviceSession
                .getOperator()
                .offline()
                .doFinally(s -> {
                    if (onDeviceUnRegister.hasDownstreams()) {
                        unregisterListener.next(childrenDeviceSession);
                    }
                    scheduleJobQueue.add(childrenDeviceSession::close);
                })
            )
            .then();
    }

    Mono<Void> tryPersistent() {
        return Flux
            .fromIterable(repository.values())
            .distinct(DeviceSession::getDeviceId)
            .filter(PersistentSession.class::isInstance)
            .cast(PersistentSession.class)
            .window(200)
            .flatMap(this::tryPersistent)
            .then();
    }

    Mono<Void> cleanSession(List<String> deviceId) {
        return sessionRepository
            .createDelete()
            .where(PersistentSessionEntity::getServerId, serverId)
            .in(PersistentSessionEntity::getDeviceId, deviceId)
            .execute()
            .onErrorResume(err -> {
                log.debug("remove session error", err);
                return Mono.empty();
            })
            .then();
    }

    Mono<Void> resumeSession(PersistentSessionEntity entity) {
        return entity
            .toSession(registry)
            .doOnNext(session -> {
                log.debug("resume session[{}]", session.getDeviceId());
                repository.put(session.getDeviceId(), session);
                repository.put(session.getId(), session);
            })
            .onErrorResume((err) -> {
                log.debug("resume session[{}] error", entity.getDeviceId(), err);
                return Mono.empty();
            })
            .then();
    }

    @Override
    public DeviceSession getSession(String clientId) {
        DeviceSession session = repository.get(clientId);

        if (session == null) {
            return null;
        }
        if (!session.isAlive()) {
            unregister(clientId);
            return null;
        }
        return session;
    }

    @Override
    public ChildrenDeviceSession getSession(String deviceId, String childrenId) {
        return Optional.ofNullable(children.get(deviceId))
                       .map(map -> map.get(childrenId))
                       .filter(ChildrenDeviceSession::isAlive)
                       .orElse(null);
    }

    @Override
    public Mono<ChildrenDeviceSession> registerChildren(String deviceId, String childrenDeviceId) {
        return Mono.defer(() -> {
            DeviceSession session = getSession(deviceId);
            if (session == null) {
                log.warn("device[{}] session not alive", deviceId);
                return Mono.empty();
            }
            if (Objects.equals(deviceId, childrenDeviceId)) {
                return Mono.empty();
            }
            return registry
                .getDevice(childrenDeviceId)
                .switchIfEmpty(Mono.fromRunnable(() -> log.debug("children device [{}] not fond in registry", childrenDeviceId)))
                .flatMap(deviceOperator -> deviceOperator
                    .online(session.getServerId().orElse(serverId), session.getId(), session
                        .getClientAddress()
                        .map(String::valueOf)
                        .orElse(null))
                    .then(deviceOperator.setConfig(DeviceConfigKey.parentGatewayId, deviceId))
                    .thenReturn(new ChildrenDeviceSession(childrenDeviceId, session, deviceOperator)))
                .doOnNext(s -> {
                    registerListener.next(s);
                    children.computeIfAbsent(deviceId, __ -> new ConcurrentHashMap<>()).put(childrenDeviceId, s);
                });
        });

    }

    @Override
    public Mono<ChildrenDeviceSession> unRegisterChildren(String deviceId, String childrenId) {

        return Mono.justOrEmpty(children.get(deviceId))
                   .flatMap(map -> Mono.justOrEmpty(map.remove(childrenId)))
                   .doOnNext(ChildrenDeviceSession::close)
                   .flatMap(session -> session.getOperator()
                                              .offline()
                                              .doFinally(s -> {
                                                  //通知
                                                  if (onDeviceUnRegister.hasDownstreams()) {
                                                      unregisterListener.next(session);
                                                  }
                                              })
                                              .thenReturn(session));
    }


    @Override
    public DeviceSession replace(DeviceSession oldSession, DeviceSession newSession) {
        DeviceSession old = repository.put(oldSession.getDeviceId(), newSession);
        if (old != null) {
            //清空sessionId不同
            if (!old.getId().equals(old.getDeviceId())) {
                repository.put(oldSession.getId(), newSession);
            }
        }
        if (newSession instanceof PersistentSession) {
            tryPersistent(Flux.just(((PersistentSession) newSession)))
                .subscribe();
        }
        return newSession;
    }

    @Override
    public DeviceSession register(DeviceSession session) {
        DeviceSession old = repository.put(session.getDeviceId(), session);
        if (old != null) {
            //清空sessionId不同
            if (!old.getId().equals(old.getDeviceId())) {
                repository.remove(old.getId());
            }
        }
        if (!session.getId().equals(session.getDeviceId())) {
            repository.put(session.getId(), session);
        }
        if (null != old) {
            if (!old.equals(session)) {
                if (old.isAlive()) {
                    //1. 可能是多个设备使用了相同的id.
                    //2. 可能是同一个设备,注销后立即上线,由于种种原因,先处理了上线后处理了注销逻辑.
                    log.warn("device[{}] session exists,disconnect old session:{}", old.getDeviceId(), session);
                    //加入关闭连接队列
                    scheduleJobQueue.add(old::close);
                }
            } else {
                return old;
            }
        } else {
            //本地计数
            getCounter(session.getTransport()).increment();
        }

        //注册中心上线
        session.getOperator()
               .online(session.getServerId().orElse(serverId), session.getId(), session
                   .getClientAddress()
                   .map(String::valueOf)
                   .orElse(null))
               .doFinally(s -> {
                   //通知
                   if (onDeviceRegister.hasDownstreams()) {
                       registerListener.next(session);
                   }
               })
               .subscribe();

        return old;
    }

    @Override
    public Flux<DeviceSession> onRegister() {
        return onDeviceRegister;
    }

    @Override
    public Flux<DeviceSession> onUnRegister() {
        return onDeviceUnRegister;
    }

    @Override
    public Flux<DeviceSession> getAllSession() {
        return Flux
            .fromIterable(repository.values())
            .distinct(DeviceSession::getDeviceId);
    }

    @Override
    public boolean sessionIsAlive(String deviceId) {
        return getSession(deviceId) != null
            ||
            children.values()
                    .stream()
                    .anyMatch(r -> {
                        DeviceSession session = r.get(deviceId);
                        return session != null && session.isAlive();
                    });
    }

    @Override
    public DeviceSession unregister(String idOrDeviceId) {
        return unregister(idOrDeviceId, true);
    }

    LongAdder getCounter(Transport transport) {
        return transportCounter
            .computeIfAbsent(transport.getId(), _transport -> new LongAdder());
    }

    DeviceSession unregister(String idOrDeviceId, boolean notify) {
        DeviceSession session = repository.remove(idOrDeviceId);

        if (null != session) {
            if (!session.getId().equals(session.getDeviceId())) {
                repository.remove(session.getId().equals(idOrDeviceId) ? session.getDeviceId() : session.getId());
            }
            //本地计数
            getCounter(session.getTransport()).decrement();
            if (notify) {
                //通知
                unregisterSession.next(session);
            }
            //加入关闭连接队列
            scheduleJobQueue.add(session::close);
        }
        return session;
    }

    public Mono<Void> tryPersistent(Flux<PersistentSession> sessions) {
        if (sessionRepository == null) {
            return Mono.empty();
        }
        return sessions
            .flatMap(session -> PersistentSessionEntity.from(serverId, session, registry))
            .distinct(PersistentSessionEntity::getId)
            .as(sessionRepository::save)
            .onErrorResume(err -> {
                log.warn("persistent session error", err);
                return Mono.empty();
            })
            .then();
    }

}